import torch
import torch.nn as nn


class Autoencoder(nn.Module):
    """
    该模型基于 U-Net 结构的变体设计，以增强自编码器 (Autoencoder) 在图像重建任务中的能力。
    主要特点：
    1. **跳跃连接 (Skip Connections)**：在编码器与解码器之间建立信息传输通道，避免深度学习模型信息丢失，提高重建细节。
    2. **适应任意输入尺寸**：模型支持任意 `H×W` 形状的输入，只要其高度和宽度能被 4 整除。
    3. **轻量级设计**：相比传统 U-Net，减少了一些计算量，使得模型更加高效，适合 CIFAR-10 等中等规模数据集。
    """

    def __init__(self):
        super().__init__()

        # -------------------- 编码器部分 --------------------
        # 第一层编码，不改变输入尺寸，仅提取特征
        self.enc1 = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, padding=1),  # 3通道（RGB） → 64通道
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=3, padding=1),  # 进一步提取特征
            nn.ReLU()
        )

        # 第二层编码，使用池化层下采样，降低特征图尺寸，提高抽象信息表达
        self.enc2 = nn.Sequential(
            nn.MaxPool2d(kernel_size=2),  # 空间尺寸缩小一半
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv2d(128, 128, kernel_size=3, padding=1),
            nn.ReLU()
        )

        # 第三层编码，进一步压缩特征表示，增强对高层语义特征的提取
        self.enc3 = nn.Sequential(
            nn.MaxPool2d(kernel_size=2),  # 再次缩小一半
            nn.Conv2d(128, 256, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv2d(256, 256, kernel_size=3, padding=1),
            nn.ReLU()
        )

        # -------------------- 解码器部分 --------------------
        # 第一层解码，使用反卷积恢复空间分辨率
        self.dec3 = nn.Sequential(
            nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2, output_padding=1),  # 4x4 → 8x8
            nn.ReLU()
        )

        # 第二层解码，继续恢复原始尺寸
        self.dec2 = nn.Sequential(
            nn.ConvTranspose2d(256, 64, kernel_size=2, stride=2, output_padding=1),  # 8x8 → 16x16
            nn.ReLU()
        )

        # 输出层，恢复到原始通道数（3通道RGB），保持输入尺寸
        self.out_layer = nn.Sequential(
            nn.Conv2d(128, 3, kernel_size=1),  # 1x1卷积用于调整通道
            nn.Sigmoid()  # 归一化到 [0,1]，适用于图像数据
        )

    def forward(self, x):
        """
        前向传播过程，包括编码和解码阶段。
        使用跳跃连接（Skip Connections）以减少信息损失，提高重建质量。

        参数：
            x (Tensor): 输入图像张量，形状为 [batch, 3, H, W]。

        返回：
            out (Tensor): 重建后的图像张量，形状与输入一致。
        """
        # 编码部分
        x1 = self.enc1(x)  # 保持原始尺寸
        x2 = self.enc2(x1)  # 下采样 H/2
        x3 = self.enc3(x2)  # 下采样 H/4

        # 解码部分（带跳跃连接）
        y = self.dec3(x3)  # H/4 → H/2
        y = torch.cat([y, x2], dim=1)  # 跳跃连接，融合细节信息

        y = self.dec2(y)  # H/2 → H
        y = torch.cat([y, x1], dim=1)  # 跳跃连接，进一步融合细节

        out = self.out_layer(y)  # 最终输出，与输入尺寸一致
        return out
